import os
from langchain.tools import BaseTool
from pydantic import BaseModel,Field,ConfigDict
from typing import Any,Type
from langchain_chroma import Chroma
from model.my_chat_model import ChatModel
from sentence_transformers import CrossEncoder
import torch

#缓存模型实例(避免每次初始化)
_reranker_model = None

def get_reranker_model():
    global _reranker_model
    if _reranker_model is None:
        print('正在加载 BAAI/bge-reranker-base模型')
        _reranker_model=CrossEncoder(
            "BAAI/bge-reranker-base",
            max_length=512,#根据文档长度自定义调整
            device= "cuda" if torch.cuda.is_available() else "cpu", #自动选择设备
        )
        print(f"模型已加载，运行设备：{_reranker_model.device}")
    return _reranker_model
#定义输入参数的数据模型类型，... 表示必填
class ChromaInput(BaseModel):
    question:str = Field(...,description="企业背景查询问题")

#定义工具类
class ChromaTool(BaseTool):
    # 定义模型，是否允许输入参数，输入参数的数据模型类
    model_config = ConfigDict(arbitrary_types_allowed=True)
    def __init__(self,**kwargs:Any):
        super().__init__(
            name = "get_chroma_tool",
            description="主要用于查询chroma数据库的企业背景信息,用户必须输入的参数是问题",
            **kwargs
        )

    #定义工具参数
    args_schema : Type[BaseModel] = ChromaInput

    def _run(self,question:str):
        # 1. 获取 embedding 模型（用于 Chroma 检索）
        chat = ChatModel()
        embedding_model = chat.get_embedding_model()
        # 服务器环境的路径
        # db = os.path.join("chroma_db")
        # 本地测试地址
        db = "../chroma_db"
        # print(f"db={db}")

        # 2. 加载 Chroma 向量库
        store = Chroma(
            persist_directory=db,
            embedding_function=embedding_model,
            collection_name="companies_context",
            collection_metadata={"hnsw:space":"cosine"}
        )
        # 3. 初步检索 top-10 文档
        # search_type="similarity" 根据文档的余玄相似度进行查询，k表示返回的文档数量
        res=store.as_retriever(search_type="similarity",search_kwargs={"k":10})
        #检索
        docs=res.invoke(question)
        print(f"docs:{docs}")
        print(f"初步检索到 {len(docs)}个文档")
        if not docs:
            return "未检索到相关企业背景信息"

        # 4. 准备 rerank 输入：[query, doc_text] 对
        doc_texts = [doc.page_content for doc in docs]
        pairs = [[question,text] for text in doc_texts]
        # 5. 使用本地 BGE Reranker 模型进行重排序
        try:
            model = get_reranker_model()
            with torch.no_grad():
                scores=model.predict(pairs)
                print(f"scores:{scores}")
        except Exception as e:
            print(f"BGE Reranker 推理失败 ：{e}")
            #降级：返回原始 top-5
            reranked_docs = docs[:5]
        else:
            #6.按得分排序，取top-1
            socred_docs = sorted(zip(scores,docs),key=lambda x:x[0],reverse=True)
            reranked_docs = [doc for _,doc in socred_docs[:1]]

        # 7. 拼接最终上下文和来源
        context_text = "\n\n".join([doc.page_content.strip() for doc in reranked_docs])

        #去重来源文件
        source_texts = list(set(doc.metadata.get("source","未知") for doc in reranked_docs))
        source_text = "\n".join(source_texts)
        whole_text = f"{context_text}\n\n数据来源:\n{source_text}"
        return whole_text

if __name__=="__main__":
    tool = ChromaTool()
    tool._run(question="请查询一下公司2的背景信息")


